#include "itkImageRegistrationMethod.h"
#include "itkMeanSquaresImageToImageMetric.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkRegularStepGradientDescentOptimizer.h"
#include "itkImage.h"
#include <itkWindowedSincInterpolateImageFunction.h>

#include "itkCenteredRigid2DTransform.h"

#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"

#include "itkResampleImageFilter.h"
#include "itkSubtractImageFilter.h"
#include "itkRescaleIntensityImageFilter.h"

#include "itkPluginUtilities.h"

#include "registerSlicesCLP.h"
#include <string>
#include <sstream>
#include <iostream>

//  The following section of code implements a Command observer
//  that will monitor the evolution of the registration process.
//
#include "itkCommand.h"
class CommandIterationUpdate : public itk::Command 
{
public:
  typedef  CommandIterationUpdate   Self;
  typedef  itk::Command             Superclass;
  typedef itk::SmartPointer<Self>  Pointer;
  itkNewMacro( Self );
protected:
  CommandIterationUpdate() {};
public:
  typedef itk::RegularStepGradientDescentOptimizer     OptimizerType;
  typedef   const OptimizerType   *    OptimizerPointer;

  void Execute(itk::Object *caller, const itk::EventObject & event)
    {
      Execute( (const itk::Object *)caller, event);
    }

  void Execute(const itk::Object * object, const itk::EventObject & event)
    {
      OptimizerPointer optimizer = 
        dynamic_cast< OptimizerPointer >( object );
      if( ! itk::IterationEvent().CheckEvent( &event ) )
        {
        return;
        }
      std::cout << optimizer->GetCurrentIteration() << "   ";
      std::cout << optimizer->GetValue() << "   ";
      std::cout << optimizer->GetCurrentPosition() << std::endl;
    }
};



template <class T>
std::string to_string(T t, std::ios_base & (*f)(std::ios_base&), int iLZ)
{
  std::ostringstream oss;
  oss << std::setfill('0') << std::setw(iLZ) << f << t;
  return oss.str();
}

template< class PixelType >
int DoIt( int argc, const char *argv[], PixelType )
{
  
  PARSE_ARGS;

  // for now use nrrd as desired suffixes
  // later we will get this automatically from the input file name

  const    unsigned int    Dimension = 2;

  // this defines the pixel type
  //typedef  unsigned short PixelType;

  typedef itk::Image< PixelType, Dimension >  FixedImageType;
  typedef itk::Image< PixelType, Dimension >  MovingImageType;

  typedef itk::CenteredRigid2DTransform< double > TransformType;

  typedef TransformType::MatrixType MatrixType;
  typedef TransformType::OutputVectorType VectorType;

  typedef itk::RegularStepGradientDescentOptimizer       OptimizerType;
  typedef itk::MeanSquaresImageToImageMetric< 
                                    FixedImageType, 
                                    MovingImageType >    MetricType;
 
  typedef itk:: LinearInterpolateImageFunction< 
                                    MovingImageType,
                                    double          >
  InterpolatorType;

  typedef itk::ImageRegistrationMethod< 
                                    FixedImageType, 
                                    MovingImageType >    RegistrationType;

  typename MetricType::Pointer         metric        = MetricType::New();
  typename OptimizerType::Pointer      optimizer     = OptimizerType::New();
  typename InterpolatorType::Pointer   interpolator  = InterpolatorType::New();
  typename RegistrationType::Pointer   registration  = RegistrationType::New();
  
  registration->SetMetric(        metric        );
  registration->SetOptimizer(     optimizer     );
  registration->SetInterpolator(  interpolator  );

  typename TransformType::Pointer  transform = TransformType::New();
  registration->SetTransform( transform );

  typedef itk::ImageFileReader< FixedImageType  > FixedImageReaderType;
  typedef itk::ImageFileReader< MovingImageType > MovingImageReaderType;

  typename FixedImageReaderType::Pointer  fixedImageReader  = FixedImageReaderType::New();
  typename MovingImageReaderType::Pointer movingImageReader = MovingImageReaderType::New();

  // need to loop through all the images

  // first define the initial rotation matrix and translation vector

  MatrixType accumulatedTransformMatrix;
  accumulatedTransformMatrix.SetIdentity();  // set to indentity

  VectorType accumulatedTransformVector;
  accumulatedTransformVector.Fill( 0 );  // set to zero displacement
  
  typedef typename FixedImageType::SpacingType    SpacingType;
  typedef typename FixedImageType::PointType      OriginType;
  typedef typename FixedImageType::RegionType     RegionType;
  typedef typename FixedImageType::SizeType       SizeType;

  CommandIterationUpdate::Pointer observer = CommandIterationUpdate::New();
  optimizer->AddObserver( itk::IterationEvent(), observer );

  // first write out the initial image that everything will be aligned to

  typedef itk::ImageFileWriter< FixedImageType >  WriterFixedType;
    
  typename WriterFixedType::Pointer      writer =  WriterFixedType::New();

  std::string inputExtension = inputVolume.substr( inputVolume.find_last_of(".") );
  std::string outputExtension = outputVolume.substr( outputVolume.find_last_of(".") );
  std::string inputBaseName = inputVolume.substr(0, inputVolume.find_last_of(".") );
  std::string outputBaseName = outputVolume.substr(0, outputVolume.find_last_of(".") );

  std::cout << "inputExtension = " << inputExtension << std::endl;
  std::cout << "outputExtension = " << outputExtension << std::endl;
  std::cout << "inputBaseName = " << inputBaseName << std::endl;
  std::cout << "outputBaseName = " << outputBaseName << std::endl;

  std::string fixedFileName = inputBaseName;
  fixedFileName += to_string<int>(iFrom, std::dec, iLZ);
  fixedFileName += inputExtension;
  
  fixedImageReader->SetFileName( fixedFileName );
  fixedImageReader->Update();

  std::string outputFileName = outputBaseName;
  outputFileName += to_string<int>(iFrom, std::dec, iLZ);
  outputFileName += outputExtension;

  writer->SetFileName( outputFileName );

  writer->SetInput( fixedImageReader->GetOutput()   );

  try
    {
    writer->Update();
    }
  catch( itk::ExceptionObject & excp )
    { 
    std::cerr << "ExceptionObject while writing the resampled image !" << std::endl; 
    std::cerr << excp << std::endl; 
    return EXIT_FAILURE;
    } 

  for ( int iI=iFrom; iI<iTo; iI++ )
    {
    std::string fixedFileName = inputBaseName;
    fixedFileName += to_string<int>(iI, std::dec, iLZ);
    fixedFileName += inputExtension;

    std::string movingFileName = inputBaseName;
    movingFileName += to_string<int>(iI+1, std::dec, iLZ);
    movingFileName += inputExtension;

    std::cout << "Computing registration between images " << iI << "(fixed) and " << iI+1 << "(moving)" << std::endl;

    std::cout << "fixedFileName = " << fixedFileName << std::endl;
    std::cout << "movingFileName = " << movingFileName << std::endl;

    
    fixedImageReader->SetFileName( fixedFileName );
    movingImageReader->SetFileName( movingFileName );

    fixedImageReader->Update();

    registration->SetFixedImage(    fixedImageReader->GetOutput()    );
    registration->SetMovingImage(   movingImageReader->GetOutput()   );

    registration->SetFixedImageRegion( 
      fixedImageReader->GetOutput()->GetBufferedRegion() );    

    fixedImageReader->Update();
    movingImageReader->Update();

    typename FixedImageType::Pointer fixedImage = fixedImageReader->GetOutput();

    const SpacingType fixedSpacing = fixedImage->GetSpacing();
    const OriginType  fixedOrigin  = fixedImage->GetOrigin();
    const RegionType  fixedRegion  = fixedImage->GetLargestPossibleRegion(); 
    const SizeType    fixedSize    = fixedRegion.GetSize();

    typename TransformType::InputPointType centerFixed;
  
    centerFixed[0] = fixedOrigin[0] + fixedSpacing[0] * fixedSize[0] / 2.0;
    centerFixed[1] = fixedOrigin[1] + fixedSpacing[1] * fixedSize[1] / 2.0;

    typename MovingImageType::Pointer movingImage = movingImageReader->GetOutput();
    
    const SpacingType movingSpacing = movingImage->GetSpacing();
    const OriginType  movingOrigin  = movingImage->GetOrigin();
    const RegionType  movingRegion  = movingImage->GetLargestPossibleRegion();
    const SizeType    movingSize    = movingRegion.GetSize();
    
    typename TransformType::InputPointType centerMoving;
    
    centerMoving[0] = movingOrigin[0] + movingSpacing[0] * movingSize[0] / 2.0;
    centerMoving[1] = movingOrigin[1] + movingSpacing[1] * movingSize[1] / 2.0;

    transform->SetCenter( centerFixed );
    transform->SetTranslation( centerMoving - centerFixed );

    transform->SetAngle( 0.0 );

    registration->SetInitialTransformParameters( transform->GetParameters() );

    // set the optimization parameters

    typedef OptimizerType::ScalesType       OptimizerScalesType;
    OptimizerScalesType optimizerScales( transform->GetNumberOfParameters() );
    const double translationScale = 1.0 / 1000.0;
    
    optimizerScales[0] = 1.0;
    optimizerScales[1] = translationScale;
    optimizerScales[2] = translationScale;
    optimizerScales[3] = translationScale;
    optimizerScales[4] = translationScale;

    optimizer->SetScales( optimizerScales );

    double initialStepLength = 0.05;

    optimizer->SetRelaxationFactor( 0.5 );
    optimizer->SetMaximumStepLength( initialStepLength ); 
    optimizer->SetMinimumStepLength( 0.001 );
    optimizer->SetNumberOfIterations( 200 );

    try 
      { 
      registration->StartRegistration(); 
      } 
    catch( itk::ExceptionObject & err ) 
      { 
      std::cerr << "ExceptionObject caught !" << std::endl; 
      std::cerr << err << std::endl; 
      return EXIT_FAILURE;
      } 

    typename OptimizerType::ParametersType finalParameters = 
      registration->GetLastTransformParameters();

    const double finalAngle           = finalParameters[0];
    const double finalRotationCenterX = finalParameters[1];
    const double finalRotationCenterY = finalParameters[2];
    const double finalTranslationX    = finalParameters[3];
    const double finalTranslationY    = finalParameters[4];
    
    const unsigned int numberOfIterations = optimizer->GetCurrentIteration();
    
    const double bestValue = optimizer->GetValue();


    // Print out results
    //
    const double finalAngleInDegrees = finalAngle * 45.0 / atan(1.0);

    std::cout << "Result = " << std::endl;
    std::cout << " Angle (radians)   = " << finalAngle  << std::endl;
    std::cout << " Angle (degrees)   = " << finalAngleInDegrees  << std::endl;
    std::cout << " Center X      = " << finalRotationCenterX  << std::endl;
    std::cout << " Center Y      = " << finalRotationCenterY  << std::endl;
    std::cout << " Translation X = " << finalTranslationX  << std::endl;
    std::cout << " Translation Y = " << finalTranslationY  << std::endl;
    std::cout << " Iterations    = " << numberOfIterations << std::endl;
    std::cout << " Metric value  = " << bestValue          << std::endl;

    // now resample, but make sure that we use the accumulated
    // transformations (to relate everything back to the first image)

    typedef itk::ResampleImageFilter< 
      MovingImageType, 
      FixedImageType >    ResampleFilterType;

    typename TransformType::Pointer finalTransform = TransformType::New();

    finalTransform->SetParameters( finalParameters );

    std::cout << "matrix = " << std::endl << finalTransform->GetMatrix() << std::endl;
    std::cout << "translation = " << std::endl << finalTransform->GetTranslation() << std::endl;

    // incorporate the previously accumulated transforms

    MatrixType newAccumulatedTransformMatrix = finalTransform->GetMatrix()*accumulatedTransformMatrix;
    VectorType newAccumulatedTransformVector = finalTransform->GetMatrix()*accumulatedTransformVector + finalTransform->GetTranslation();

    accumulatedTransformMatrix = newAccumulatedTransformMatrix;
    accumulatedTransformVector = newAccumulatedTransformVector;

    finalTransform->SetTranslation( accumulatedTransformVector );
    finalTransform->SetMatrix( accumulatedTransformMatrix );

    typedef itk::WindowedSincInterpolateImageFunction<MovingImageType, 4 > ResampleInterpolatorType;
    typename ResampleInterpolatorType::Pointer   resampleInterpolator  = ResampleInterpolatorType::New();

    typename ResampleFilterType::Pointer resample = ResampleFilterType::New();

    resample->SetTransform( finalTransform );
    resample->SetInput( movingImageReader->GetOutput() );
    resample->SetInterpolator( resampleInterpolator );
    
    resample->SetSize(    fixedImage->GetLargestPossibleRegion().GetSize() );
    resample->SetOutputOrigin(  fixedImage->GetOrigin() );
    resample->SetOutputSpacing( fixedImage->GetSpacing() );
    resample->SetDefaultPixelValue( 0 );
    
    typename WriterFixedType::Pointer      writer =  WriterFixedType::New();

    std::string outputFileName = outputBaseName;
    outputFileName += to_string<int>(iI+1, std::dec, iLZ);
    outputFileName += outputExtension;

    writer->SetFileName( outputFileName );

    writer->SetInput( resample->GetOutput()   );

    try
      {
      writer->Update();
      }
    catch( itk::ExceptionObject & excp )
      { 
      std::cerr << "ExceptionObject while writing the resampled image !" << std::endl; 
      std::cerr << excp << std::endl; 
      return EXIT_FAILURE;
      } 

    }

  return EXIT_SUCCESS;
}

int main ( int argc, const char *argv[] )
{
  PARSE_ARGS;
  itk::ImageIOBase::IOPixelType pixelType;
  itk::ImageIOBase::IOComponentType componentType;

  std::string inputExtension = inputVolume.substr( inputVolume.find_last_of(".") );
  std::string inputBaseName = inputVolume.substr(0, inputVolume.find_last_of(".") );

  std::string firstFile = inputBaseName;
  firstFile += to_string<int>(iFrom, std::dec, iLZ );
  firstFile += inputExtension;

  try
    {
    itk::GetImageType (firstFile, pixelType, componentType);

    // This filter handles all types
    
    switch (componentType)
      {
      case itk::ImageIOBase::UCHAR: 
	std::cout << "Image type: UCHAR" << std::endl;
	return DoIt( argc, argv, static_cast<unsigned char>(0) ); break;
      case itk::ImageIOBase::CHAR:
	std::cout << "Image type: CHAR" << std::endl;
	return DoIt( argc, argv, static_cast<char>(0) ); break;
      case itk::ImageIOBase::USHORT:
	std::cout << "Image type: USHORT" << std::endl;
  	return DoIt( argc, argv, static_cast<unsigned short>(0) ); break;
      case itk::ImageIOBase::SHORT:
	std::cout << "Image type: SHORT" << std::endl;
  	return DoIt( argc, argv, static_cast<short>(0) ); break;
      case itk::ImageIOBase::UINT:
	std::cout << "Image type: UINT" << std::endl;
  	return DoIt( argc, argv, static_cast<unsigned int>(0) ); break;
      case itk::ImageIOBase::INT:
	std::cout << "Image type: INT" << std::endl;
  	return DoIt( argc, argv, static_cast<int>(0) ); break;
      case itk::ImageIOBase::ULONG:
	std::cout << "Image type: ULONG" << std::endl;
  	return DoIt( argc, argv, static_cast<unsigned long>(0) ); break;
      case itk::ImageIOBase::LONG:
	std::cout << "Image type: LONG" << std::endl;
  	return DoIt( argc, argv, static_cast<long>(0) ); break;
      case itk::ImageIOBase::FLOAT:
	std::cout << "Image type: FLOAT" << std::endl;
  	return DoIt( argc, argv, static_cast<float>(0) ); break;
      case itk::ImageIOBase::DOUBLE:
	std::cout << "Image type: DOUBLE" << std::endl;
  	return DoIt( argc, argv, static_cast<double>(0) ); break;
      default: std::cerr << "Unknown image type. Abort." << std::endl;
	return EXIT_FAILURE;
      }

    }
  catch( itk::ExceptionObject &excep)
    {
    std::cerr << argv[0] << ": exception caught !" << std::endl;
    std::cerr << excep << std::endl;
    return EXIT_FAILURE;
    }

  return EXIT_SUCCESS;

}

